import shutil
import sys
import time
import xml.etree.ElementTree as ET
import numpy as np
from collections import defaultdict
from datetime import datetime, timedelta
from orbitP.script import config
import orekit
from orekit.pyhelpers import setup_orekit_curdir
import pandas as pd
from org.hipparchus.geometry.euclidean.threed import Vector3D
from org.orekit.frames import FramesFactory
from org.orekit.time import TimeScalesFactory, AbsoluteDate
from org.orekit.utils import IERSConventions, PVCoordinates
from tqdm import tqdm
import os

dataFINDir = "../dataset/dataFIN/"
dataRAPDir = "../dataset/dataRAP/"
dataSULTDir = "../dataset/dataSULT/"
dataObsDir = "../dataset/dataCSV/dataObs/"
dataPrdDir = "../dataset/dataCSV/dataPrd/"
dataFinalDir = "../dataset/dataCSV/dataFinal/"
dataRapidDir = "../dataset/dataCSV/dataRapid/"

beginDate = datetime.strptime("2023-06-01 00:00:00", "%Y-%m-%d %H:%M:%S")
endDate   = datetime.strptime("2025-08-10 23:59:59", "%Y-%m-%d %H:%M:%S")
dataType = 2

def get_file_paths(directory):
    file_paths = []
    for root, dirs, files in os.walk(directory):
        for file in files:
            file_path = os.path.join(root, file)
            file_paths.append(os.path.normpath(file_path))
    return file_paths

def yeardoy2date(year, doy, hs):
    year = int(year)
    doy = int(doy)
    hr = int(hs[:2])
    mite = int(hs[2:])
    return datetime(year, 1, 1) + timedelta(days=doy - 1, hours=hr, minutes=mite)


def ecef_to_eci(pList, vList, tList):
    # """
    #     将 ECEF 坐标批量转换为 ECI 坐标。
    #
    #     参数:
    #         pList (np.ndarray): ECEF 位置数组，形状为 (N, 3)。
    #         vList (np.ndarray): ECEF 速度数组，形状为 (N, 3)。
    #         tList (list): 时间列表，长度为 N。
    #
    #     返回:
    #         np.ndarray: 转换后的 ECI 位置数组，形状为 (N, 3)。
    #         np.ndarray: 转换后的 ECI 速度数组，形状为 (N, 3)。
    #     """
    # 初始化时间系统和参考框架
    utc = TimeScalesFactory.getUTC()

    # 使用EME2000作为ECI框架
    eci_frame = FramesFactory.getEME2000()
    ecef_frame = FramesFactory.getITRF(IERSConventions.IERS_2010, True)

    N = len(tList)
    rECI = np.zeros((N, 3), dtype=np.float64)
    vECI = np.zeros((N, 3), dtype=np.float64)

    for i in range(N):
        if isinstance(tList[i], np.datetime64):
            timeNow = tList[i].astype(datetime.datetime)
        else:
            timeNow = tList[i]  # 已经是 Python datetime
        date = AbsoluteDate(
            timeNow.year, timeNow.month, timeNow.day,
            timeNow.hour, timeNow.minute,
            timeNow.second + timeNow.microsecond / 1e6,
            utc
        )

        r_vec = Vector3D(float(pList[i, 0]), float(pList[i, 1]), float(pList[i, 2]))
        v_vec = Vector3D(float(vList[i, 0]), float(vList[i, 1]), float(vList[i, 2]))
        pv = PVCoordinates(r_vec, v_vec)

        transform = ecef_frame.getTransformTo(eci_frame, date)
        pv_eci = transform.transformPVCoordinates(pv)

        rECI[i, :] = pv_eci.getPosition().toArray()
        vECI[i, :] = pv_eci.getVelocity().toArray()
    return rECI, vECI

def insertRapidORFinalV(dataRes,dataSLTime):
    dataNewRes = {} #存储新的数据
    for dateName, data in dataRes.items():
        dateList = dateName.split("_")
        date = datetime.strptime(dateList[0],"%Y-%m-%d") #日期
        sat_id = dateList[1] #卫星编号
        if date > endDate or date < beginDate:
            continue
        dataNew = [] #存储当前天的时刻数据
        flag = True
        for idx, dataNow in enumerate(data): #枚举当天的时刻
            dataTimeNow = dataNow[0]
            preTime = dataTimeNow - timedelta(minutes=15) #前一个时刻
            preName = preTime.strftime("%Y-%m-%d %H:%M:%S") + "_" + sat_id
            if preName not in dataSLTime.keys():
                flag=False
                break
            preP = dataSLTime[preName] #前一个时刻的位置

            nxtTime = dataTimeNow + timedelta(minutes=15) #后一个时刻
            nxtName = nxtTime.strftime("%Y-%m-%d %H:%M:%S") + "_" + sat_id
            if nxtName not in dataSLTime.keys():
                flag=False
                break
            nxtP = dataSLTime[nxtName] #后一个时刻的位置

            dt_prev = (dataTimeNow - preTime).total_seconds()
            dt_next = (nxtTime - dataTimeNow).total_seconds()
            nowV = [(nxtP[i] - preP[i]) / (dt_prev + dt_next) for i in range(len(nxtP))] #计算速度
            newRow = dataNow[:4] + nowV + [dataNow[-1]] # 时间 位置 速度 卫星钟差
            dataNew.append(newRow)
        if flag == True:
            dataNewRes[dateName] = dataNew
    return dataNewRes

def insertObsANDPrdV(dataObs_ori, dataPrd_ori, dataSTime,useNear=True):
    dataObs = {} #存储新的观测段和预测段
    dataPrd = {}
    lack = 0
    for fileName, data in dataObs_ori.items():
        obs = np.array(data)
        prd = np.array(dataPrd_ori[fileName])
        fileNameList = fileName.split("_")
        fileDate = datetime.strptime(fileNameList[0],"%Y-%m-%d %H-%M-%S")
        sat_id = fileNameList[-1]
        preFileDate = fileDate - timedelta(days=1) # 获取前后完整不重叠的观测段和预测段
        sufFileDate = fileDate + timedelta(days=1)
        preFileName = preFileDate.strftime("%Y-%m-%d %H-%M-%S") + "_" + sat_id
        sufFileName = sufFileDate.strftime("%Y-%m-%d %H-%M-%S") + "_" + sat_id
        if fileDate < beginDate or fileDate > endDate:
            continue
        if preFileName not in dataSTime or sufFileName not in dataSTime:
            lack+=1
            continue
        concatData = np.vstack((obs, prd)) #合并观测段和预测段进行计算速度
        if useNear == False:
            rangeList = range(1,concatData.shape[0]-1)
        else:
            rangeList = range(concatData.shape[0])
        for idx in rangeList:
            dataTimeNow = concatData[idx][0]
            preTime = dataTimeNow - timedelta(minutes=15)
            preName = preTime.strftime("%Y-%m-%d %H-%M-%S") + "_" + sat_id
            nxtTime = dataTimeNow + timedelta(minutes=15)
            nxtName = nxtTime.strftime("%Y-%m-%d %H-%M-%S") + "_" + sat_id
            if idx == 0: #最开始
                preP = dataSTime[preFileName][preName]
                nxtP = concatData[idx+1][1:4]
            elif idx == concatData.shape[0] - 1: #最末尾
                preP = concatData[idx-1][1:4]
                nxtP = dataSTime[fileName][nxtName]
            else:
                preP = concatData[idx-1][1:4]
                nxtP = concatData[idx+1][1:4]

            dt_prev = (dataTimeNow - preTime).total_seconds()
            dt_next = (nxtTime - dataTimeNow).total_seconds()
            nowV = np.array([(nxtP[i] - preP[i]) / (dt_prev + dt_next) for i in range(len(nxtP))]).reshape(1, -1)
            newRow = np.concatenate((concatData[idx:idx+1,:config.axis+1], nowV, concatData[idx:idx+1,-1:]), axis=1)
            if idx < obs.shape[0]:
                dataObs.setdefault(fileName, []).append(newRow)
            else:
                dataPrd.setdefault(fileName, []).append(newRow)

    print(f"lack: {lack}")
    for key in dataObs:
        dataObs[key] = np.vstack(dataObs[key])
    for key in dataPrd:
        dataPrd[key] = np.vstack(dataPrd[key])
    return dataObs, dataPrd


def get_RapidORFinal(): # 获取 快速轨道或最终轨道每天的数据
    dataRes = {} #存储每天的轨道数据
    dataSLTime = {} #存储每个时刻的数据，用于计算速度
    t_bar = tqdm(filePaths, total=len(filePaths))
    for filePath in t_bar:
        start_time = filePath.split('_')[1]
        start_time = yeardoy2date(start_time[:4], start_time[4:7],start_time[7:10]) # 开始时间
        end_time = start_time + timedelta(hours=23, minutes=59, seconds=59) #结束时间
        t_bar.set_description(f"{start_time.strftime('%Y-%m-%d')}")
        if start_time < (beginDate - timedelta(days=1)) or end_time > (endDate + timedelta(days=1)): # 不在区间内就跳过
            continue
        with open(filePath, 'r') as f:
            lines = f.readlines()
            dateTime = None
            for line in lines:
                if line.startswith('*'):  # 时间标识行
                    parts = line.split()
                    year, month, day, hour, minute, second = map(float, parts[1:7])
                    dateTime = datetime(int(year), int(month), int(day), int(hour), int(minute), int(second))

                elif line.startswith('P') and line[1:2] == "G":  # 位置数据行
                    sat_id = line[1:4]
                    if sat_id not in config.GPSList:
                        continue
                    parts = line.split()
                    x, y, z, cb = map(float, parts[1:5])  # 位置（km）
                    dateNow = [dateTime, x * 1000, y * 1000, z * 1000, cb]  # 转换为 m
                    date = dateNow[0].strftime('%Y-%m-%d')
                    fileName = date +"_"+ sat_id #每天标记是日期 + 卫星编号
                    dataName = dateTime.strftime("%Y-%m-%d %H:%M:%S") +"_"+ sat_id #时刻标记是时间 + 卫星编号
                    if dateNow[0] > end_time or dateNow[0] < start_time:
                        continue
                    dataSLTime[dataName] = [x * 1000, y * 1000, z * 1000]  # 设置映射
                    if (dateNow[0].hour * 60 + dateNow[0].minute) % 15 != 0: # 仅保留15min的时间间隔数据
                        continue
                    dataRes.setdefault(fileName, []).append(dateNow)

    dataRes = insertRapidORFinalV(dataRes,dataSLTime) #对每个时刻的速度进行计算
    t_bar = tqdm(dataRes.items(), total=len(dataRes.items()))
    for fileName, data in t_bar:
        data = np.array(data)
        data[:,1:4], data[:,4:7] = ecef_to_eci(data[:,1:4], data[:,4:7], data[:,0])
        dataRes[fileName] = data

    for fileName, data in t_bar:
        if dataType == 0:
            writePath = dataFinalDir + f"{fileName}.csv"
        else:
            writePath = dataRapidDir + f"{fileName}.csv"
        if dataType == 0:
            df = pd.DataFrame(data, columns=["date", "X", "Y", "Z", "VX", "VY", "VZ", "CB"])
            # df = pd.DataFrame(data, columns=["date", "X", "Y", "Z", "CB"])
        else:
            df = pd.DataFrame(data, columns=["date", "X", "Y", "Z", "VX", "VY", "VZ", "CB"])
            # df = pd.DataFrame(data, columns=["date", "X", "Y", "Z", "CB", "VX", "VY", "VZ"])
            # df = pd.DataFrame(data, columns=["date", "X", "Y", "Z", "CB"])
        df['date'] = pd.to_datetime(df['date'])
        df.to_csv(writePath, index=False, chunksize=50000)

def get_ObsANDPrd():
    t_bar = tqdm(filePaths, total=len(filePaths))
    dataSTime = defaultdict(dict) #存储日期-时刻的双重映射
    dataObs = {} #存储观测段
    dataPrd = {} #存储预测段
    for filePath in t_bar:
        startOb = filePath.split('_')[1]
        startOb = yeardoy2date(startOb[:4], startOb[4:7], startOb[7:10]) #观测段开始时间
        endOb = startOb + timedelta(hours=23, minutes=59, seconds=59) #观测段结束时间
        startPd = startOb + timedelta(days=1) #预测段开始时间
        endPd = endOb + timedelta(hours=23, minutes=59, seconds=59) #预测段结束时间
        # print(f"{startOb}   {endOb}")
        # print(f"{endPd}   {endPd}")
        t_bar.set_description(f"{startOb.strftime('%Y-%m-%d %H:%M:%S')}")
        if startOb < (beginDate - timedelta(days=1)) or endPd > (endDate + timedelta(days=1)):
            continue
        with open(filePath, 'r') as f:
            lines = f.readlines()
            dateTime = None
            for line in lines:
                if line.startswith('*'):  # 时间标识行
                    parts = line.split()
                    year, month, day, hour, minute, second = map(float, parts[1:7])
                    dateTime = datetime(int(year), int(month), int(day), int(hour), int(minute), int(second))
                elif line.startswith('P') and line[1:2] == "G":  # 位置数据行
                    sat_id = line[1:4]
                    if sat_id not in config.GPSList:
                        continue
                    parts = line.split()
                    x, y, z, cb = map(float, parts[1:5])  # 位置（km）
                    dateNow = [dateTime, x * 1000, y * 1000, z * 1000, cb]  # 转换为 m
                    fileName = startOb.strftime('%Y-%m-%d %H-%M-%S') + "_" + sat_id #以当前观测段开始时间作为标签
                    dataName = dateTime.strftime("%Y-%m-%d %H-%M-%S") + "_" + sat_id
                    if (dateTime.minute==45) or (dateTime.minute==0): #45为开始的前一个时刻，0为结束的后一个时刻
                        dataSTime[fileName][dataName] = [x * 1000, y * 1000, z * 1000]
                    if dateNow[0] > endPd or dateNow[0] < startOb or (
                            dateNow[0].hour * 60 + dateNow[0].minute) % 15 != 0:
                        continue
                    if dateNow[0] >= startOb and dateNow[0] <= endOb: #处于观测段
                        dataObs.setdefault(fileName, []).append(dateNow)
                    elif dateNow[0] >= startPd and dateNow[0] <= endPd: #处于预测段
                        dataPrd.setdefault(fileName, []).append(dateNow)

    dataObs, dataPrd = insertObsANDPrdV(dataObs, dataPrd, dataSTime,useNear=True)
    del dataSTime

    t_bar = tqdm(dataObs.items(), total=len(dataObs.items()))
    for fileName, data in t_bar:
        data[:,1:4], data[:,4:7] = ecef_to_eci(data[:,1:4], data[:,4:7], data[:,0])
        dataObs[fileName] = data
    for fileName,data in dataObs.items():
        writePath = dataObsDir + f"{fileName}.csv"
        df = pd.DataFrame(data, columns=["date", "X", "Y", "Z", "VX", "VY", "VZ", "CB"])
        # df = pd.DataFrame(data, columns=["date", "X", "Y", "Z", "CB"])
        df['date'] = pd.to_datetime(df['date'])
        df.to_csv(writePath, index=False, chunksize=50000)

    t_bar = tqdm(dataPrd.items(), total=len(dataPrd.items()))
    for fileName, data in t_bar:
        data[:,1:4], data[:,4:7] = ecef_to_eci(data[:,1:4], data[:,4:7], data[:,0])
        dataPrd[fileName] = data
    for fileName, data in dataPrd.items():
        writePath = dataPrdDir + f"{fileName}.csv"
        df = pd.DataFrame(data, columns=["date", "X", "Y", "Z", "VX", "VY", "VZ", "CB"])
        # df = pd.DataFrame(data, columns=["date", "X", "Y", "Z", "CB"])
        df['date'] = pd.to_datetime(df['date'])
        df.to_csv(writePath, index=False, chunksize=50000)

if __name__ == "__main__":
    # 初始化 Orekit 环境
    orekit.initVM()
    config.orekitCurPath = "../tools/orekit-data-master"
    setup_orekit_curdir(config.orekitCurPath)

    if dataType == 0:
        filePaths= get_file_paths(dataFINDir)
        if os.path.exists(dataFinalDir):
            shutil.rmtree(dataFinalDir)
        if not os.path.exists(dataFinalDir):
            os.makedirs(dataFinalDir)
    elif dataType == 1:
        filePaths= get_file_paths(dataRAPDir)
        if os.path.exists(dataRapidDir):
            shutil.rmtree(dataRapidDir)
        if not os.path.exists(dataRapidDir):
            os.makedirs(dataRapidDir)
    elif dataType == 2:
        filePaths= get_file_paths(dataSULTDir)
        if os.path.exists(dataObsDir):
            shutil.rmtree(dataObsDir)
        if not os.path.exists(dataObsDir):
            os.makedirs(dataObsDir)
        if os.path.exists(dataPrdDir):
            shutil.rmtree(dataPrdDir)
        if not os.path.exists(dataPrdDir):
            os.makedirs(dataPrdDir)
    else:
        print("data not found")
        sys.exit(0)

    if dataType == 0 or dataType == 1:
        get_RapidORFinal()
    elif dataType == 2:
        get_ObsANDPrd()
    else:
        print("data not found")
        sys.exit(0)
